import streamlit as st
import pickle
import networkx as nx
import matplotlib.pyplot as plt

# 读取知识图谱
with open('knowledge_graph.gpickle', 'rb') as f:
    G = pickle.load(f)

nodes = list(G.nodes)

# 题号与知识点描述映射
node_label_map = {
    "选择1": "命题逻辑：蕴含式的成假赋值",
    "选择2": "谓词逻辑：自由变元与约束变元的区分",
    "选择3": "谓词逻辑：自由变元与约束变元的区分",
    "选择4": "关系性质：对称性判断",
    "选择5": "关系运算：复合关系",
    "选择6": "图论：度数序列可图性",
    "选择7": "图论：有向图的连通性分类",
    "选择8": "图论：生成子图的同构计数",
    "选择9": "树的性质：叶子节点计算",
    "选择10": "图论：欧拉图与哈密尔顿图的判别条件",
    "填空1": "命题逻辑：必要条件符号化",
    "填空2": "关系闭包：自反对称闭包的构造",
    "填空3": "图论：最小生成树权重",
    "填空4": "图论：握手定理应用",
    "填空5": "树的性质：顶点度数计算",
    "填空6": "图论：欧拉通路的识别",
    "计算1": "偏序集",
    "计算2": "图论",
    "证明1": "命题逻辑：自然推理系统的构造法",
    "证明2": "命题逻辑：归谬法（反证法）的应用",
    "综合应用1": "编码理论",
    "综合应用2": "图论"
}

st.title("知识点先导关系可视化查询")

# 下拉框和多选框用题号
error_node = st.selectbox("请选择要追踪的错误题目：", nodes)
correct_nodes = st.multiselect("请选择已答对的题目（可多选）：", nodes)

def find_break_nodes(G, error_node, correct_nodes):
    ancestors = nx.ancestors(G, error_node)
    start_nodes = [n for n in ancestors if G.in_degree(n) == 0]
    if not start_nodes:
        start_nodes = list(ancestors)
    break_nodes = set()
    for start in start_nodes:
        for path in nx.all_simple_paths(G, source=start, target=error_node):
            last_correct = None
            for n in path[:-1]:
                if n in correct_nodes:
                    last_correct = n
            if last_correct:
                idx = path.index(last_correct)
                if idx + 1 < len(path) and path[idx + 1] != error_node:
                    break_nodes.add(path[idx + 1])
            else:
                # 只在路径第一个节点不是error_node时才加入
                if path[0] != error_node:
                    break_nodes.add(path[0])
    # 最终只保留在先导知识点集合中的节点
    break_nodes = break_nodes & ancestors
    return break_nodes

if error_node:
    all_preds = nx.ancestors(G, error_node)
    direct_preds = list(G.predecessors(error_node))
    st.write(f"**{node_label_map[error_node]} 的直接先导知识点：** {', '.join([node_label_map[n] for n in direct_preds]) if direct_preds else '无'}")
    st.write(f"**{node_label_map[error_node]} 的所有先导知识点：** {', '.join([node_label_map[n] for n in sorted(all_preds)]) if all_preds else '无'}")

    break_nodes = find_break_nodes(G, error_node, set(correct_nodes))
    st.write(f"**推断的中断点（疑似知识断点）：** {', '.join([node_label_map[n] for n in break_nodes]) if break_nodes else '无'}")

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False

    # 可视化子图
    sub_nodes = all_preds | {error_node} | break_nodes | set(correct_nodes)
    subG = G.subgraph(sub_nodes)
    pos = nx.spring_layout(subG)
    plt.figure(figsize=(10, 7))  # 增加图表大小以容纳图例

    # 设置节点颜色
    node_colors = []
    for n in subG.nodes:
        if n in break_nodes:
            node_colors.append('red')
        elif n == error_node:
            node_colors.append('orange')
        elif n in correct_nodes:
            node_colors.append('green')
        else:
            node_colors.append('lightblue')

    # 绘制图形
    nx.draw(subG, pos, with_labels=True, labels=None, node_color=node_colors, edge_color='gray', node_size=1200, font_size=10)
    
    # 添加图例
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=15, label='知识断点'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='orange', markersize=15, label='错误题目'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='green', markersize=15, label='已答对题目'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightblue', markersize=15, label='其他知识点')
    ]
    plt.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.15, 1))
    
    # 调整布局以确保图例完全显示
    plt.tight_layout()
    st.pyplot(plt)
